#-
# ==========================================================================
# Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
#+

###############################################################################
##
## instanceShape.py
##
## Description:
##    Registers a new shape that acts like an instancer.  The new shape
##    type is called "instanceShape".
##
##    The shape will instance N copies of a shape connected via a message
##    attribute on the node.  The sample will distribute these N copies
##    in the XZ plane.
##
##    There are no output attributes for this shape.
##    The following input attributes define the type of shape to draw.
##
##       radius		   : circle radius for instance object. 
##		 instanceShape : a connection to the shape to instance 
##		 count		   : number of instances to make. 
##
##    Additionally the instancing feature demonstrated in this code
##    only works for custom shapes. Non-custom shapes will not work.
##
################################################################################

# Usage:
# import maya
# maya.cmds.loadPlugin("instanceShape.py")
# maya.cmds.loadPlugin("basicShape.py")

# basicShape = maya.cmds.createNode("spBasicShape")
# instanceShape = maya.cmds.createNode("spInstanceShape")
# maya.cmds.connectAttr( basicShape + ".message", instanceShape + ".instanceShape" )
# 

import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaRender as OpenMayaRender
import maya.OpenMayaUI as OpenMayaUI

import math
import sys

kPluginNodeTypeName = "spInstanceShape"
spInstanceShapeNodeId = OpenMaya.MTypeId(0x00080029)

glRenderer = OpenMayaRender.MHardwareRenderer.theRenderer()
glFT = glRenderer.glFunctionTable()

kLeadColor 				= 18 # green
kActiveColor			= 15 # white
kActiveAffectedColor	= 8  # purple
kDormantColor			= 4  # blue
kHiliteColor			= 17 # pale blue

kDefaultRadius = 1.0
kDefaultCount  = 10


#####################################################################
##
## Geometry class
##
class instanceGeom:
	def __init__(self): 
		self.radius = kDefaultRadius
		self.count = kDefaultCount
		self.instanceShape = None
		self.drawQueueList = [] 
		
#####################################################################
##
## Shape class - defines the non-UI part of a shape node
##
class instanceShape(OpenMayaMPx.MPxSurfaceShape):
	# class variables
	aRadius = OpenMaya.MObject()
	aCount = OpenMaya.MObject()
	aInstanceShape = OpenMaya.MObject()
	
	def __init__(self):
		OpenMayaMPx.MPxSurfaceShape.__init__(self)
		
		# geometry
		self.__myGeometry = instanceGeom()

	# override
	def postConstructor(self):
		"""
		 When instances of this node are created internally, the
		 MObject associated with the instance is not created until
		 after the constructor of this class is called. This means
		 that no member functions of MPxSurfaceShape can be called in
		 the constructor.  The postConstructor solves this
		 problem. Maya will call this function after the internal
		 object has been created.  As a general rule do all of your
		 initialization in the postConstructor.
		"""
		self.setRenderable(True)

	# override
	def getInternalValue(self, plug, datahandle):
		"""
		 Handle internal attributes.
		 In order to impose limits on our attribute values we
		 mark them internal and use the values in fGeometry instead.
		"""
		if (plug == instanceShape.aRadius):
			datahandle.setDouble(self.__myGeometry.radius)
		elif (plug == instanceShape.aCount):
			datahandle.setInt(self.__myGeometry.count)
		else:
			return OpenMayaMPx.MPxSurfaceShape.getInternalValue(self, plug, datahandle)

		return True


	# override
	def setInternalValue(self, plug, datahandle):
		"""
		 Handle internal attributes.
		 In order to impose limits on our attribute values we
		 mark them internal and use the values in fGeometry instead.
		"""

		# the minimum radius is 0
		#
		if (plug == instanceShape.aRadius):
			radius = datahandle.asDouble()
			
			if (radius < 0):
				radius = 0

			self.__myGeometry.radius = radius

		elif (plug == instanceShape.aCount):
			count = datahandle.asInt()
			if (count < 0):
				count = 0 
			self.__myGeometry.count = count
		else:
			return OpenMayaMPx.MPxSurfaceShape.setInternalValue(plug, datahandle)
		return True


	# override
	def isBounded(self):
		return True


	# override
	def boundingBox(self):
		"""
		 Returns the bounding box for the shape.
		 In this case just use the radius and height attributes
		 to determine the bounding box.
		"""
		result = OpenMaya.MBoundingBox()

		geom = self.geometry()
		
		# Include the instance shape bounding box
		if geom.instanceShape: 
			fnDag = OpenMaya.MFnDagNode( geom.instanceShape ) 
			result = fnDag.boundingBox() 

		r = geom.radius
		instanceBbox = OpenMaya.MBoundingBox( result ) 
		for c in range( geom.count ): 
			percent = float(c)/float(geom.count)
			rad = 2*math.pi * percent
			p = (r*math.cos(rad), r*math.sin(rad),0.0)
			newbbox = OpenMaya.MBoundingBox( instanceBbox )
			trans = OpenMaya.MTransformationMatrix( ) 	
			vec = OpenMaya.MVector( p[0], p[1], p[2] ) 
			trans.setTranslation( vec, OpenMaya.MSpace.kTransform ) 
			mmatrix = trans.asMatrix(); 
			newbbox.transformUsing( mmatrix ) 
			result.expand( newbbox ) 
		
		return result


	def geometry(self):
		"""
		 This function gets the values of all the attributes and
		 assigns them to the fGeometry. Calling MPlug::getValue
		 will ensure that the values are up-to-date.
		"""
		# return self.__myGeometry

		this_object = self.thisMObject()

		plug = OpenMaya.MPlug(this_object, instanceShape.aRadius)
		self.__myGeometry.radius = plug.asDouble()
		plug = OpenMaya.MPlug(this_object, instanceShape.aCount)
		self.__myGeometry.count = plug.asInt()

		plug = OpenMaya.MPlug(this_object, instanceShape.aInstanceShape)
		plugArray = OpenMaya.MPlugArray()
		plug.connectedTo( plugArray, True, False )
		if ( plugArray.length() > 0 ): 
			node = plugArray[0].node()
			dagNode = OpenMaya.MFnDagNode(node)
			path = OpenMaya.MDagPath()
			dagNode.getPath(path)
			self.__myGeometry.instanceShape = path 

		return self.__myGeometry

#####################################################################
##
## UI class	- defines the UI part of a shape node
##
class instanceShapeUI(OpenMayaMPx.MPxSurfaceShapeUI):
	# private enums
	def __init__(self):
		OpenMayaMPx.MPxSurfaceShapeUI.__init__(self)

	# override
	def getDrawRequests(self, info, objectAndActiveOnly, queue):
		"""
		 The draw data is used to pass geometry through the 
		 draw queue. The data should hold all the information
		 needed to draw the shape.
		"""
		self.geometry = None

		# Custom instancer objects can instance other custom surface 
		# shapes.  This is done by first getting the draw request 
		# data for the instancer shape 
		#
		data = OpenMayaUI.MDrawData()
		request = info.getPrototype(self)
		shapeNode = self.surfaceShape()
		path = info.multiPath()
		view = info.view() 
		# We stored the instance object via a connection on the surface
		# shape.  Retrieve that value.
		#
		geom = shapeNode.geometry()
		shadedMode = False
		mainMaterial = None 
		if request.displayStyle() == OpenMayaUI.M3dView.kGouraudShaded:
			shadedMode = True
			
		if geom.instanceShape:
			# Find the MPxSurfaceShape for the object that we are instancing
			#
			shapeUI = OpenMayaMPx.MPxSurfaceShapeUI.surfaceShapeUI(geom.instanceShape)
			if shapeUI: 	
				mainMaterial = shapeUI.material( geom.instanceShape )
				mainMaterial.evaluateMaterial( view, geom.instanceShape ) 
				r = geom.radius 
				for a in range(geom.count):
					myQueue = OpenMayaUI.MDrawRequestQueue()
					percent = float(a)/float(geom.count)
					rad = 2*math.pi * percent
					position = (r*math.cos(rad), r*math.sin(rad),0.0)
					# Construct a reference to MDrawInfo and modify it 
					# to point to the instance shape.  If we do not do 
					# this in then the call to getDrawRequests will think
					# that we are still the instancer shape and not the 
					# instance shape.  
					# 
					myinfo = OpenMayaUI.MDrawInfo( info )
					myinfo.setMultiPath( geom.instanceShape )
					shapeUI.getDrawRequests( myinfo,
											 objectAndActiveOnly, myQueue )
 					geom.drawQueueList.append( (myQueue, position) )

		info.setMultiPath( path )
		# Finally we must supply a material back to the drawing code. 
		# We attempt to use the instance shape material; however, if 
		# that fails then we fall back to the default material 
		# 
 		if shadedMode:
			defaultMaterial = OpenMayaUI.MMaterial.defaultMaterial()
 			if not mainMaterial:
 				mainMaterial = defaultMaterial 
			try: 
				request.setMaterial( mainMaterial )
			except:
				request.setMaterial( defaultMaterial ) 

			
		self.getDrawData(geom, data)			
		request.setDrawData(data)

		self.geometry = geom
		queue.add(request)

	# override
	def draw(self, request, view):
		"""
		 From the given draw request, get the draw data and determine
		 which basic shape to draw and with what values.
		"""
		
		data = request.drawData()
		shapeNode = self.surfaceShape()
		geom = self.geometry
		glFT.glMatrixMode( OpenMayaRender.MGL_MODELVIEW )
		if geom.instanceShape: 
			shapeUI = OpenMayaMPx.MPxSurfaceShapeUI.surfaceShapeUI(geom.instanceShape)
			for (queue, pos) in geom.drawQueueList:
				glFT.glPushMatrix();
				glFT.glTranslatef( pos[0], pos[1], pos[2] )
				while not queue.isEmpty():
					request = queue.remove()
					shapeUI.draw( request, view )
				glFT.glPopMatrix()

		# Draw a shell area that shows where the instances are being
		# drawn. This is nice to have if we don't have any instance
		# shapes connected to this plugin. 
		# 
		glFT.glPushAttrib( OpenMayaRender.MGL_ALL_ATTRIB_BITS )
		glFT.glPolygonMode(OpenMayaRender.MGL_FRONT_AND_BACK,
						   OpenMayaRender.MGL_LINE)
		glFT.glBegin(OpenMayaRender.MGL_QUADS)
		glFT.glVertex3f(-1*(geom.radius), -1*(geom.radius), 0.0)
		glFT.glNormal3f(0, 0, 1.0)
			
		glFT.glVertex3f(-1*(geom.radius), (geom.radius), 0.0)
		glFT.glNormal3f(0, 0, 1.0)
			
		glFT.glVertex3f((geom.radius), (geom.radius), 0.0)
		glFT.glNormal3f(0, 0, 1.0)
			
		glFT.glVertex3f((geom.radius), -1*(geom.radius), 0.0)
		glFT.glNormal3f(0, 0, 1.0)
		glFT.glEnd()
		glFT.glPopAttrib( )

	# override
	def select(self, selectInfo, selectionList, worldSpaceSelectPts):
		"""
		 Select function. Gets called when the bbox for the object is
		 selected.  This function just selects the object without
		 doing any intersection tests.
		"""
		
		priorityMask = OpenMaya.MSelectionMask(OpenMaya.MSelectionMask.kSelectObjectsMask)
		item = OpenMaya.MSelectionList()
		item.add(selectInfo.selectPath())
		xformedPt = OpenMaya.MPoint()
		selectInfo.addSelection(item, xformedPt, selectionList,
								 worldSpaceSelectPts, priorityMask, False)
		return True




#####################################################################

def nodeCreator():
	return OpenMayaMPx.asMPxPtr( instanceShape() )


def uiCreator():
	return OpenMayaMPx.asMPxPtr( instanceShapeUI() )


def nodeInitializer():
	# utility func for numeric attrs
	def setOptions(attr):
		attr.setHidden(False)
		attr.setKeyable(True)
		attr.setInternal(True)

	messageAttr = OpenMaya.MFnMessageAttribute() 
	numericAttr = OpenMaya.MFnNumericAttribute()

	instanceShape.aInstanceShape = messageAttr.create("instanceShape", "is")
	instanceShape.addAttribute(instanceShape.aInstanceShape)
	
	instanceShape.aRadius = numericAttr.create("radius", "r", OpenMaya.MFnNumericData.kDouble, kDefaultRadius)
	setOptions(numericAttr)
	instanceShape.addAttribute(instanceShape.aRadius)

	instanceShape.aCount = numericAttr.create("count", "ct", OpenMaya.MFnNumericData.kInt, kDefaultCount)
	setOptions(numericAttr)
	instanceShape.addAttribute(instanceShape.aCount)
	
# initialize the script plug-in
def initializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject, "Autodesk", "2011", "Any")
	try:
		mplugin.registerShape( kPluginNodeTypeName, spInstanceShapeNodeId,
								nodeCreator, nodeInitializer, uiCreator )
	except:
		sys.stderr.write( "Failed to register node: %s" % kPluginNodeTypeName )
		raise


# uninitialize the script plug-in
def uninitializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject)
	try:
		mplugin.deregisterNode( spInstanceShapeNodeId )
	except:
		sys.stderr.write( "Failed to deregister node: %s" % kPluginNodeTypeName )
		raise
